package com.custom;


import com.custom.source.ClinkSplitEnumerator;
import lombok.val;
import lombok.var;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;

import java.io.IOException;

/**
 * 具体的分片枚举器
 */
public class MySplitEnumerator extends ClinkSplitEnumerator<MySplit, MyEnumCheckpoint> {
    private long to;
    private long from;

    public MySplitEnumerator(SplitEnumeratorContext<MySplit> enumContext, long from, long to) {
        super(enumContext);
        this.from = from;
        this.to = to;
    }

    public MySplitEnumerator(SplitEnumeratorContext<MySplit> enumContext, MyEnumCheckpoint checkpoint) {
        super(enumContext, checkpoint);
    }

    /**
     * 我选择在这里进行拆分
     */
    @Override
    public void start() {
        if (!pendingSplits.isEmpty()) {
            return;
        }

        val parallelism = context.currentParallelism();
        if (parallelism < 1) {
            throw new IllegalArgumentException("parallelism must be greater than 1");
        }
        if (parallelism == 1) {
            pendingSplits.add(new MySplit(0, from, to));
        }

        if (parallelism > 1) {
            val length = to - from;
            val step = length % parallelism == 0 ? length / parallelism : length / parallelism + 1;

            var tmp = from;
            for (int i = 0; i < parallelism; i++) {
                pendingSplits.add(new MySplit(i, tmp, tmp + step));
                tmp = tmp + step;
            }
        }
    }

    /**
     * 这个我也不知道啥意思
     */
    @Override
    public void addReader(int subtaskId) {
    }

    /**
     * 释放一些资源,如果有需要的话
     */
    @Override
    public void close() throws IOException {
    }
}
